"""
数据库管理模块
负责SQLite数据库的初始化和操作
"""
import sqlite3
import os
from typing import List, Dict, Any, Optional
from .config import config


class Database:
    """数据库管理类"""
    
    def __init__(self):
        self.db_path = config.get('database.path', 'data/database/tiny_qa.db')
        self._init_database()
    
    def _init_database(self):
        """初始化数据库"""
        os.makedirs(os.path.dirname(self.db_path), exist_ok=True)
        
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            
            # 创建文档表
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS documents (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    file_path TEXT UNIQUE NOT NULL,
                    file_name TEXT NOT NULL,
                    file_type TEXT NOT NULL,
                    file_size INTEGER,
                    content TEXT,
                    processed BOOLEAN DEFAULT FALSE,
                    chunk_count INTEGER DEFAULT 0,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            # 创建文档块表
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS document_chunks (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    document_id INTEGER NOT NULL,
                    chunk_index INTEGER NOT NULL,
                    content TEXT NOT NULL,
                    vector_id TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY (document_id) REFERENCES documents (id)
                )
            ''')
            
            # 创建问答历史表
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS qa_history (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    question TEXT NOT NULL,
                    answer TEXT NOT NULL,
                    sources TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            # 创建系统配置表
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS system_config (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    key TEXT UNIQUE NOT NULL,
                    value TEXT,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            conn.commit()
    
    def get_connection(self):
        """获取数据库连接"""
        return sqlite3.connect(self.db_path)
    
    def execute_query(self, query: str, params: tuple = ()) -> List[Dict[str, Any]]:
        """执行查询"""
        with self.get_connection() as conn:
            conn.row_factory = sqlite3.Row
            cursor = conn.cursor()
            cursor.execute(query, params)
            return [dict(row) for row in cursor.fetchall()]
    
    def execute_update(self, query: str, params: tuple = ()) -> int:
        """执行更新操作"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute(query, params)
            conn.commit()
            return cursor.rowcount
    
    def insert_document(self, file_path: str, file_name: str, file_type: str, 
                       file_size: int, content: str = None) -> int:
        """插入文档记录"""
        query = '''
            INSERT OR REPLACE INTO documents 
            (file_path, file_name, file_type, file_size, content, updated_at)
            VALUES (?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
        '''
        self.execute_update(query, (file_path, file_name, file_type, file_size, content))
        
        # 获取插入的ID
        result = self.execute_query(
            "SELECT id FROM documents WHERE file_path = ?", 
            (file_path,)
        )
        return result[0]['id'] if result else None
    
    def get_document(self, file_path: str) -> Optional[Dict[str, Any]]:
        """获取文档信息"""
        result = self.execute_query(
            "SELECT * FROM documents WHERE file_path = ?", 
            (file_path,)
        )
        return result[0] if result else None
    
    def get_all_documents(self) -> List[Dict[str, Any]]:
        """获取所有文档"""
        return self.execute_query("SELECT * FROM documents ORDER BY created_at DESC")
    
    def update_document_processed(self, file_path: str, processed: bool, chunk_count: int = 0):
        """更新文档处理状态"""
        query = '''
            UPDATE documents 
            SET processed = ?, chunk_count = ?, updated_at = CURRENT_TIMESTAMP
            WHERE file_path = ?
        '''
        self.execute_update(query, (processed, chunk_count, file_path))
    
    def insert_chunk(self, document_id: int, chunk_index: int, content: str, vector_id: str = None) -> int:
        """插入文档块"""
        query = '''
            INSERT INTO document_chunks 
            (document_id, chunk_index, content, vector_id)
            VALUES (?, ?, ?, ?)
        '''
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute(query, (document_id, chunk_index, content, vector_id))
            conn.commit()
            return cursor.lastrowid
    
    def get_chunks_by_document(self, document_id: int) -> List[Dict[str, Any]]:
        """获取文档的所有块"""
        return self.execute_query(
            "SELECT * FROM document_chunks WHERE document_id = ? ORDER BY chunk_index",
            (document_id,)
        )
    
    def insert_qa_history(self, question: str, answer: str, sources: str = None):
        """插入问答历史"""
        query = '''
            INSERT INTO qa_history (question, answer, sources)
            VALUES (?, ?, ?)
        '''
        self.execute_update(query, (question, answer, sources))
    
    def get_qa_history(self, limit: int = 50) -> List[Dict[str, Any]]:
        """获取问答历史"""
        return self.execute_query(
            "SELECT * FROM qa_history ORDER BY created_at DESC LIMIT ?",
            (limit,)
        )
    
    def set_config(self, key: str, value: str):
        """设置系统配置"""
        query = '''
            INSERT OR REPLACE INTO system_config (key, value, updated_at)
            VALUES (?, ?, CURRENT_TIMESTAMP)
        '''
        self.execute_update(query, (key, value))
    
    def get_config(self, key: str) -> Optional[str]:
        """获取系统配置"""
        result = self.execute_query(
            "SELECT value FROM system_config WHERE key = ?",
            (key,)
        )
        return result[0]['value'] if result else None


# 全局数据库实例
db = Database() 